#!/bin/bash
#
# Unlock TCG Opal 2-compliant disks and reboot to continue booting the 'real' operating system
#
# This script executes the final stage of a TCG Opal pre-boot authentication (PBA) boot.
# It is expected that this script executes on a volatile system running entirely on RAM file systems.
# To avoid delays, this script will perform a hard reset or power-off instead of a regular
# system shutdown.

shopt -s nullglob extglob  # Enable pattern matching extensions required for 'opal-functions.sh'
source /usr/share/rear/lib/opal-functions.sh
[[ -f /.OPAL_PBA_SETTINGS.sh ]] && source /.OPAL_PBA_SETTINGS.sh
[[ ${#OPAL_PBA_TKNPATH[@]} -gt 0 ]] && source /usr/share/rear/lib/authtoken-functions.sh
[[ -n "$OPAL_PBA_TPMNVINDEX" ]] && source /usr/share/rear/lib/tpmpwd-functions.sh


function use_plymouth() {
    # returns 0 if plymouth is to be used.
    type -p plymouth &>/dev/null && plymouth --ping
}

function quit_plymouth() {
    # quits plymouth if in use.
    use_plymouth && plymouth quit
}

function enter_plymouth_shutdown_mode() {
    # puts plymouth into shutdown mode, if in use.
    use_plymouth && plymouth change-mode --shutdown
}

function display_message() {
    local message="${1:?}"
    (use_plymouth && plymouth display-message --text="$message") || echo -e "\n$message"
}

function hide_message() {
    #local message="${1:?}"
    #use_plymouth && plymouth hide-message --text="$message"
    # hide-message doesn't work, at least with two-step plugin
    use_plymouth && plymouth display-message --text=""
    return 0
}

function ask_for_password() {
    local prompt="${1:?}"  # a colon will be appended implicitly
    local allow_empty="$2"
    # asks for a password, setting the variable $password.
    local msg_req="Please enter a non-empty password"

    while true; do
        if ! password="$(use_plymouth && plymouth ask-for-password --prompt="$prompt")"; then
            echo ""
            read -r -s -p "$prompt: " password 2>&1
            echo ""
        fi
        [[ -n "$password" || -n "$allow_empty" ]] && break || display_message "$msg_req"
    done
    hide_message "$msg_req"

    [[ -n "$password" ]] && is_debug_password "$password" && emergency_shell
    return 0
}

function ask_question() {
    local prompt="${1:?}"
    local allow_empty="$2"
    # asks question, setting the variable $answer.
    local msg_req="Answer is required"

    while true; do
        if ! answer="$(use_plymouth && plymouth ask-question --prompt="$prompt")"; then
            echo ""
            read -r -p "$prompt " answer 2>&1
            echo ""
        fi
        [[ -n "$answer" || -n "$allow_empty" ]] && break || display_message "$msg_req"
    done
    hide_message "$msg_req"
    return 0
}

function is_debug_password() {
    local password="${1:?}"
    # returns true (0) if $password matches the configured debug password.
    # Will return false (1) if there is no openssl executable to verify the password hash.

    [[ -z "$OPAL_PBA_DEBUG_PASSWORD" ]] && return 1  # debug password not configured
    type -p openssl &>/dev/null || return 1  # openssl not available

    # Extract hash algorithm and salt from the configured hash signature
    local method salt _
    IFS='\$' read -r _ method salt _ <<<"$OPAL_PBA_DEBUG_PASSWORD"

    # Compute signature of the password to verify
    password_signature="$(openssl passwd -"$method" -salt "$salt" "$password")"

    # Return the verification result
    [[ "$password_signature" == "$OPAL_PBA_DEBUG_PASSWORD" ]]
}

function is_secboot() {
    #return: 0=enabled, 1=disabled, 2=undef/unsupported

    # 8be4df61-93ca-11d2-aa0d-00e098032b8c is the global ns GUID
    local efi_var="/sys/firmware/efi/efivars/SecureBoot-8be4df61-93ca-11d2-aa0d-00e098032b8c"
    local tmp

    if [ -r "$efi_var" ]; then
        tmp="$(od -An -tu1 -j4 $efi_var)"
        [[ "${tmp//[[:space:]]/}" == "1" ]] && return 0 || return 1
    else
        return 2
    fi
}

error_log="/error-log.$$"
exec 2> "$error_log"

function emergency_shell() {
    # drop into an emergency shell, then shutdown.

    exec 2>&1

    display_message "Entering emergency shell..."
    sleep 3
    quit_plymouth

    local history_file="/.bash_history.$$" rc_file="/.bashrc.$$"

    cat > "$history_file" << '--EOF--'
sedutil-cli --help |& less  # help on low-level administration of Opal-compliant disks
sedutil-cli --scan  # scan for Opal-compliant disks
sedutil-cli --query /dev/sda
journalctl  # show system startup log
exit
--EOF--

    cat > "$rc_file" << '--EOF--'
export PS1="OPAL PBA> "
rear() {
    echo "ERROR: You cannot run rear from within the OPAL PBA." >&2
    return 1
}
cd /
--EOF--

    cat > /etc/motd << '--EOF--'

This is the OPAL PBA emergency shell.

See history for useful commands. Exit the shell to shut down the system.

--EOF--

    if [[ -s "$error_log" ]]; then
        {
            echo "The following errors occurred when executing $0:"
            cat "$error_log"
            echo ""
        } >> /etc/motd
    fi

    HISTFILE="$history_file" bash --rcfile "$rc_file"

    exit  # Terminate the shell, initiating a shutdown.
}

function stop_error_handling() {
    trap - EXIT
}

function instant_reboot() {
    enter_plymouth_shutdown_mode
    stop_error_handling

    # Force immediate hardware reboot via Magic SysRq key
    echo 1 > /proc/sys/kernel/sysrq
    echo b > /proc/sysrq-trigger

    # Fallback if the previous method did not work
    sleep 1
    reboot --force
}

function instant_poweroff() {
    enter_plymouth_shutdown_mode
    stop_error_handling

    # Force immediate hardware poweroff via Magic SysRq key
    echo 1 > /proc/sys/kernel/sysrq
    echo o > /proc/sysrq-trigger

    # Fallback if the previous method did not work
    sleep 1
    poweroff --force
}

function handle_termination() {
    # handle script termination because of an error, interrupt, or a user exiting the debug shell.

    display_message "Shutting down..."
    sleep 3
    instant_poweroff
}


trap handle_termination EXIT


if use_plymouth; then
    # Initialize boot splash screen animation if available
    plymouth update-root-fs --read-write
else
    # Clear screen if running without plymouth boot animation and if 'clear' is available
    type -p clear &>/dev/null && clear
fi


# Minimal system setup
# TODO: split system setup scripts into PBA and rescue categories to protect against script renaming
for system_setup_script in 00-functions.sh 10-console-setup.sh 40-start-udev-or-load-modules.sh; do
    source "/etc/scripts/system-setup.d/$system_setup_script"
done

# Find TCG Opal 2-compliant disks
devices=( $(opal_devices) )
declare -i device_count=${#devices[@]}
[[ -n "$OPAL_PBA_DEBUG_DEVICE_COUNT" ]] && device_count="$OPAL_PBA_DEBUG_DEVICE_COUNT"
if (( device_count == 0 )); then
    display_message "Could not detect TCG Opal 2-compliant disks."
    echo "Could not detect TCG Opal 2-compliant disks." >&2
    sleep 3
    exit
fi

# Query TCG Opal 2-compliant disks to determine the maximum number of authentication attempts
declare -i max_authentications=5  # self-imposed limit to begin with
for device in "${devices[@]}"; do
    device_max_authentications="$(opal_device_max_authentications "$device")"
    if (( device_max_authentications > 0 && device_max_authentications < max_authentications )); then
        # Limit authentication attempts to the lowest number supported by any disk
        max_authentications=$device_max_authentications
    fi
done

if [[ "$OPAL_PBA_UNLOCK_MODE" == "permanent" ]]; then
    # Unlock devices permanently. A workaround mode for systems requiring a power cycle after transient unlocking.
    unlock_command=opal_device_deactivate_locking
    result_state_message="unlocked permanently"
else
    # Unlock devices transiently until powering off. This is the normal operation for protection at rest.
    unlock_command=opal_device_unlock
    result_state_message="unlocked"
fi

# Ask for a password, unlock TCG Opal 2-compliant disks, reboot if successful
if (( device_count == 1 )); then
    password_prompt="Enter password to unlock disk"
    unsuccessful_unlock_response="Could not unlock the disk."
else
    password_prompt="Enter password to unlock disks"
    unsuccessful_unlock_response="Could not unlock any of $device_count disks."
fi

# OPAL_PBA_SBWARN is either "pseudo-bool" or string warning message to show
# Newer Plymouth with two-step plugin has similar check built-in and can show a warning-image
if ! is_false "$OPAL_PBA_SBWARN" && ! is_secboot; then
    sb_fail_status=$?
    if is_true "$OPAL_PBA_SBWARN"; then
        # Plymouth itself is responsible for displaying a warning in this case
        # Showing default warning msg if Plymouth not available
        use_plymouth || display_message "WARNING! SecureBoot is DISABLED or UNAVAILABLE"
    else
        display_message "${OPAL_PBA_SBWARN:-WARNING! SecureBoot is DISABLED or UNAVAILABLE}"
    fi
fi

authtkn_bind2dev=""
is_true "$OPAL_PBA_TKNBIND" && authtkn_bind2dev="/dev/disk/by-partlabel/${OPAL_PBA_GPT_PARTITION_NAME// /\\x20}"
authtkn_2famkr=$'\005' #ASCII non-printable 'enquiry'

authtkn_allowed=""
[[ ${#OPAL_PBA_TKNPATH[@]} -gt 0 ]] && authtkn_allowed="Y"

authtkn_stderr="/dev/null"
if is_true "$OPAL_PBA_TKNDBG"; then
    echo "AuthToken bind: [$OPAL_PBA_TKNBIND] -> $authtkn_bind2dev"
    authtkn_stderr="/.authtkn.stderr.$$"
fi

tpmpwd_allowed=""
[[ ${#OPAL_PBA_TKNPATH[@]} -eq 0 && -n "$OPAL_PBA_TPMNVINDEX" ]] && tpmpwd_allowed="Y"

tpmpwd_stderr="/dev/null"
if is_false "$OPAL_PBA_TKNDBG" && is_true "$OPAL_PBA_TPMDBG"; then
    echo "TPM bind: [$OPAL_PBA_TKNBIND] -> $authtkn_bind2dev"
    tpmpwd_stderr="/.tpmpwd.stderr.$$"
fi

declare -i attempt=0
declare -i attempts2fa=0
declare -i unlocked_device_count=0

while (( unlocked_device_count == 0 && attempt < max_authentications )); do
    if [ -n "$authtkn_allowed" ] && [ -z "$password" ]; then
        # just booted or user has explicitly requested AuthToken-load by supplying an empty password
        for authtkn_path in "${OPAL_PBA_TKNPATH[@]}"; do
            blkdev_wait "$authtkn_path" 4 1 2>>$authtkn_stderr && break
            authtkn_path=""
        done

        if [ -n "$authtkn_path" ]; then # AuthToken container is present
            is_true "$OPAL_PBA_TKNDBG" && echo "AuthToken container[$authtkn_path] present as $(blkdev_model $authtkn_path 2>>$authtkn_stderr)"
            at_password=$(authtkn_load "$authtkn_path" "$OPAL_PBA_TKNOFFSET" "$OPAL_PBA_TKNKEY" "$authtkn_bind2dev" 2>>$authtkn_stderr)
            if [ $? -eq 0 ]; then #loaded & decrypted successfully
                is_true "$OPAL_PBA_TKNDBG" && echo "AuthToken successfully loaded"
                if [ "${at_password:0:${#authtkn_2famkr}}" == "$authtkn_2famkr" ]; then # AuthToken is 2FA protected
                    is_true "$OPAL_PBA_TKNDBG" && echo "AuthToken is 2FA-protected"
                    attempts2fa=${OPAL_PBA_TKN2FAMAXTRIES:-3}
                    while (( attempts2fa > 0 )); do
                        is_true "$OPAL_PBA_TKNDBG" && echo "AuthToken 2FA attempt $attempts2fa"
                        ask_for_password "Please, type your 2FA password"
                        password=$(decrypt_base64 "${at_password:${#authtkn_2famkr}}" "$password" 2>>$authtkn_stderr)
                        [[ $? -eq 0 ]] && break || (( attempts2fa -= 1 ))
                    done

                    if [ -z "$password" ]; then # user failed to provide correct 2FA password
                        is_true "$OPAL_PBA_TKNDBG" && echo "AuthToken 2FA failed, disabling AT"
                        authtkn_allowed="" # disabling AT
                        if is_true "$OPAL_PBA_TKN2FAFAILWIPE"; then
                            is_true "$OPAL_PBA_TKNDBG" && echo "AuthToken 2FA-FAIL-WIPE enabled, wiping token"
                            authtkn_wipe "$authtkn_path" "$OPAL_PBA_TKNOFFSET" >>$authtkn_stderr 2>&1
                        fi
                    fi
                elif [ -z "$sb_fail_status" ]; then # SecBoot is active or user explicitly doesn't care, allowing unattended boot
                    password="$at_password"
                else # SecBoot disabled or unavailable, ensure user is aware of that
                    ask_question "SecureBoot is DISABLED or UNAVAILABLE! Continue anyway [Y/N]?"
                    is_true "$answer" && password="$at_password" || instant_reboot
                fi
            fi
        fi
    elif [ -n "$tpmpwd_allowed" ] && [ -z "$password" ]; then
        at_password=$(tpmpwd_load "$OPAL_PBA_TPMNVINDEX" 2>>$tpmpwd_stderr)
        if [ $? -eq 0 ]; then #loaded & decrypted successfully
            is_true "$OPAL_PBA_TPMDBG" && echo "TPM password successfully loaded"
            if [ -z "$sb_fail_status" ]; then # SecBoot is active or user explicitly doesn't care, allowing unattended boot
                password="$at_password"
            else # SecBoot disabled or unavailable, ensure user is aware of that
                ask_question "SecureBoot is DISABLED or UNAVAILABLE! Continue anyway [Y/N]?"
                is_true "$answer" && password="$at_password" || instant_reboot
            fi
        fi
    else
        # this is to cover multiple paths with the single [ -z "$password" ] block below
        password=""
    fi

    if [ -z "$password" ]; then
        # AuthToken not-configured/runtime-disabled/otherwise-failed or prev-iter $password was rejected
        at_password="" #this var also acts as indicator that password is AT-derived
        ask_for_password "$password_prompt" "$authtkn_allowed"
    fi

    if [ -n "$password" ]; then
        # Success in this case is achieved if at least one device can be unlocked.
        # If other devices require different passwords for unlocking, we assume
        # that this is intentional and will be dealt with by other means.
        attempt+=1
        if is_true "$OPAL_PBA_TKNDBG" || is_true "$OPAL_PBA_TPMDBG"; then
            # acceptable typed-in only
            [[ "$password" == "dbgpass" ]] && [[ -z "$at_password" ]] && unlocked_device_count+=1
            # acceptable both AT-sourced or typed-in
            [[ "$password" == "dbgpasstkn" ]] && unlocked_device_count+=1
        else
            for device in "${devices[@]}"; do
                "$unlock_command" "$device" "$password" >/dev/null && unlocked_device_count+=1
            done
        fi
        if [ -n "$at_password" ] && (( unlocked_device_count == 0 )); then
            # AuthToken-derived password unlocked nothing -> stop considering AT as unlock option
            authtkn_allowed="" # disabling AT
            tpm_allowed=""
            is_true "$OPAL_PBA_TKNDBG" && echo "AuthToken-derived password REJECTED"
            is_true "$OPAL_PBA_TPMDBG" && echo "TPM-derived password REJECTED"
        fi
    fi
done

if (( unlocked_device_count > 0 )); then
    if [ -z "$at_password" ] && [ ${#OPAL_PBA_TKNPATH[@]} -gt 0 ]; then
        # password is user-supplied & valid
        opal_password="$password"
        for authtkn_path in "${OPAL_PBA_TKNPATH[@]}"; do
            blkdev_wait "$authtkn_path" 4 1 2>>$authtkn_stderr || continue
            # AuthToken container present

            is_true "$OPAL_PBA_TKNDBG" && echo "Offering to store AuthToken onto $authtkn_path at offset ${OPAL_PBA_TKNOFFSET}"
            ask_question "Want to use $(blkdev_model $authtkn_path 2>>$authtkn_stderr) as AuthToken media [Y/N]?"
            if is_true "$answer"; then
                ask_for_password "Please, type optional 2FA password (or leave empty)" "Y" #allow empty
                if [ -n "$password" ]; then
                    is_true "$OPAL_PBA_TKNDBG" && echo "Encrypting AuthToken with 2FA password..."
                    if at_password=$(encrypt_base64 "$opal_password" "$password" 2>>$authtkn_stderr); then
                        at_password="${authtkn_2famkr}${at_password}"
                    else
                        display_message "2FA-encrypt failed, AuthToken NOT generated"
                    fi
                else
                    at_password="$opal_password"
                fi
                if [ -n "$at_password" ]; then
                    is_true "$OPAL_PBA_TKNDBG" && echo "AuthToken generated, storing..."
                    if authtkn_store "$authtkn_path" "$OPAL_PBA_TKNOFFSET" "$at_password" "$OPAL_PBA_TKNKEY" "$authtkn_bind2dev" >>$authtkn_stderr 2>&1; then
                        display_message "Success! $(blkdev_model $authtkn_path 2>>$authtkn_stderr) is now your AuthToken media for this PBA"
                    else
                        display_message "Failed to generate/store AuthToken onto $(blkdev_model $authtkn_path 2>>$authtkn_stderr)"
                    fi
                fi
                sleep 3
            fi
        done
    elif [ -z "$at_password" ] && [ -n "$OPAL_PBA_TPMNVINDEX" ]; then
        # password is user-supplied & valid
        opal_password="$password"
        if [ -e /dev/tpm0 ]; then
            is_true "$OPAL_PBA_TPMDBG" && echo "Offering to store password onto TPM"
            ask_question "Want to store password on TPM [Y/N]?"
            if is_true "$answer"; then
                at_password="$opal_password"
                is_true "$OPAL_PBA_TPMDBG" && echo "Storing..."
                if tpmpwd_store "$OPAL_PBA_TPMNVINDEX" "$at_password" >>$tpmpwd_stderr 2>&1; then
                    display_message "Success! password is now stored in TPM for this PBA"
                else
                    display_message "Failed to store password on TPM 2>>$tpmpwd_stderr)"
                fi
            fi
        fi
    fi

    if ! is_true "$OPAL_PBA_NOSUCCESSMSG"; then
        if (( device_count == 1 && unlocked_device_count == 1 )); then
            display_message "Disk $result_state_message, rebooting..."
        else
            display_message "$unlocked_device_count of $device_count disks $result_state_message, rebooting..."
        fi
        sleep 1
    fi
    is_true "$OPAL_PBA_TKNDBG" || is_true "$OPAL_PBA_TPMDBG" && emergency_shell || instant_reboot
else
    display_message "$unsuccessful_unlock_response"
    # If finally unsuccessful, power off.
    # This is required as TCG Opal 2-compliant disks will refuse further authentication attempts before being reset.
    display_message "Powering off after $attempt unsuccessful attempts..."
    sleep 3
    is_true "$OPAL_PBA_TKNDBG" || is_true "$OPAL_PBA_TPMDBG" && emergency_shell || instant_poweroff
fi
